package snapshotter

import (
	"context"

	snapshotsapi "github.com/containerd/containerd/api/services/snapshots/v1"
	"github.com/containerd/containerd/api/types"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/snapshots"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type Service struct {
	snapshotter snapshots.Snapshotter
	snapshotsapi.UnimplementedSnapshotsServer
}

func NewService(sn snapshots.Snapshotter) *Service {
	return &Service{
		snapshotter: sn,
	}
}

func (s *Service) Prepare(ctx context.Context, req *snapshotsapi.PrepareSnapshotRequest) (*snapshotsapi.PrepareSnapshotResponse, error) {
	mounts, err := s.snapshotter.Prepare(ctx, req.Key, req.Parent, snapshots.WithLabels(req.Labels))
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.PrepareSnapshotResponse{
		Mounts: toProtoMounts(mounts),
	}, nil
}

func (s *Service) View(ctx context.Context, req *snapshotsapi.ViewSnapshotRequest) (*snapshotsapi.ViewSnapshotResponse, error) {
	mounts, err := s.snapshotter.View(ctx, req.Key, req.Parent, snapshots.WithLabels(req.Labels))
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.ViewSnapshotResponse{
		Mounts: toProtoMounts(mounts),
	}, nil
}

func (s *Service) Mounts(ctx context.Context, req *snapshotsapi.MountsRequest) (*snapshotsapi.MountsResponse, error) {
	mounts, err := s.snapshotter.Mounts(ctx, req.Key)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.MountsResponse{
		Mounts: toProtoMounts(mounts),
	}, nil
}

func (s *Service) Commit(ctx context.Context, req *snapshotsapi.CommitSnapshotRequest) (*emptypb.Empty, error) {
	if err := s.snapshotter.Commit(ctx, req.Name, req.Key, snapshots.WithLabels(req.Labels)); err != nil {
		return nil, toGRPCError(err)
	}

	return &emptypb.Empty{}, nil
}

func (s *Service) Remove(ctx context.Context, req *snapshotsapi.RemoveSnapshotRequest) (*emptypb.Empty, error) {
	if err := s.snapshotter.Remove(ctx, req.Key); err != nil {
		return nil, toGRPCError(err)
	}

	return &emptypb.Empty{}, nil
}

func (s *Service) Stat(ctx context.Context, req *snapshotsapi.StatSnapshotRequest) (*snapshotsapi.StatSnapshotResponse, error) {
	info, err := s.snapshotter.Stat(ctx, req.Key)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.StatSnapshotResponse{
		Info: toProtoInfo(info),
	}, nil
}

func (s *Service) Update(ctx context.Context, req *snapshotsapi.UpdateSnapshotRequest) (*snapshotsapi.UpdateSnapshotResponse, error) {
	info, err := s.snapshotter.Update(ctx, fromProtoInfo(req.Info), req.UpdateMask.GetPaths()...)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.UpdateSnapshotResponse{
		Info: toProtoInfo(info),
	}, nil
}

func (s *Service) List(req *snapshotsapi.ListSnapshotsRequest, srv snapshotsapi.Snapshots_ListServer) error {
	return s.snapshotter.Walk(srv.Context(), func(ctx context.Context, info snapshots.Info) error {
		return srv.Send(&snapshotsapi.ListSnapshotsResponse{
			Info: []*snapshotsapi.Info{toProtoInfo(info)},
		})
	}, req.Filters...)
}

func (s *Service) Usage(ctx context.Context, req *snapshotsapi.UsageRequest) (*snapshotsapi.UsageResponse, error) {
	usage, err := s.snapshotter.Usage(ctx, req.Key)
	if err != nil {
		return nil, toGRPCError(err)
	}

	return &snapshotsapi.UsageResponse{
		Size:   usage.Size,
		Inodes: usage.Inodes,
	}, nil
}

func (s *Service) Cleanup(ctx context.Context, req *snapshotsapi.CleanupRequest) (*emptypb.Empty, error) {
	// Optional: implement cleanup logic
	return &emptypb.Empty{}, nil
}

// Helper functions

func toProtoMounts(mounts []mount.Mount) []*types.Mount {
	var result []*types.Mount
	for _, m := range mounts {
		result = append(result, &types.Mount{
			Type:    m.Type,
			Source:  m.Source,
			Options: m.Options,
		})
	}
	return result
}

func toProtoInfo(info snapshots.Info) *snapshotsapi.Info {
	return &snapshotsapi.Info{
		Name:      info.Name,
		Parent:    info.Parent,
		Kind:      snapshotsapi.Kind(info.Kind),
		CreatedAt: timestamppb.New(info.Created),
		UpdatedAt: timestamppb.New(info.Updated),
		Labels:    info.Labels,
	}
}

func fromProtoInfo(info *snapshotsapi.Info) snapshots.Info {
	return snapshots.Info{
		Name:    info.Name,
		Parent:  info.Parent,
		Kind:    snapshots.Kind(info.Kind),
		Created: info.CreatedAt.AsTime(),
		Updated: info.UpdatedAt.AsTime(),
		Labels:  info.Labels,
	}
}

func toGRPCError(err error) error {
	if err == nil {
		return nil
	}

	switch {
	case errdefs.IsNotFound(err):
		return status.Error(codes.NotFound, err.Error())
	case errdefs.IsAlreadyExists(err):
		return status.Error(codes.AlreadyExists, err.Error())
	case errdefs.IsInvalidArgument(err):
		return status.Error(codes.InvalidArgument, err.Error())
	case errdefs.IsFailedPrecondition(err):
		return status.Error(codes.FailedPrecondition, err.Error())
	default:
		return status.Error(codes.Unknown, err.Error())
	}
}
